"""
Minimal example script for converting a dataset to LeRobot format.

We use the Libero dataset (stored in RLDS) for this example, but it can be easily
modified for any other data you have saved in a custom format.
You can download the raw Libero datasets from https://huggingface.co/datasets/openvla/modified_libero_rlds
By default the resulting dataset will get saved to the $LEROBOT_HOME directory.
Copy the dataset to path/to/lerobot_modify/data/modified_libero_rlds to run this script

Usage:
python examples/12_rlds2lerobot_dataset.py

Note: to run the script, you need to install tensorflow_datasets:
`pip install tensorflow tensorflow_datasets`
"""
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
import shutil

from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
import tensorflow_datasets as tfds

RAW_DATASET_NAMES = [
    # "libero_10_no_noops",
    # "libero_goal_no_noops",
    "libero_object_no_noops",
    # "libero_spatial_no_noops",
]  # For simplicity we will combine multiple Libero datasets into one training dataset


def lerobot_builder(data_dir: str, *, push_to_hub: bool = False):
    # Clean up any existing dataset in the output directory
    output_path = os.path.join(os.path.dirname(__file__), "../data/")

    # Create LeRobot dataset, define features to store
    # OpenPi assumes that proprio is stored in `state` and actions in `action`
    # LeRobot assumes that dtype of image data is `image`
    dataset = LeRobotDataset.create(
        repo_id='modified_libero_lerobot',
        root=os.path.join(output_path, "modified_libero_lerobot"),
        robot_type="panda",
        fps=10,
        features={
            "observation.image.image": {
                "dtype": 'video',
                "shape": (256, 256, 3),
                "names": ["height", "width", "channels"],
            },
            "observation.image.wrist_image": {
                "dtype": 'video',
                "shape": (256, 256, 3),
                "names": ["height", "width", "channels"],
            },
            "observation.state": {
                "dtype": "float32",
                "shape": (8,),
                "names": ["state"],
            },
            "action": {
                "dtype": "float32",
                "shape": (7,),
                "names": ["action"],
            },
        },
        image_writer_threads=10,
        image_writer_processes=5,
    )

    # Loop over raw Libero datasets and write episodes to the LeRobot dataset
    # You can modify this for your own data format
    for raw_dataset_name in RAW_DATASET_NAMES:
        raw_dataset = tfds.load(raw_dataset_name, data_dir=data_dir, split="train")
        for episode in raw_dataset:
            for step in episode["steps"].as_numpy_iterator():
                dataset.add_frame(
                    {
                        "observation.image.image": step["observation"]["image"],
                        "observation.image.wrist_image": step["observation"]["wrist_image"],
                        "observation.state": step["observation"]["state"],
                        "action": step["action"],
                        "task": step["language_instruction"].decode(),
                    }
                )
            dataset.save_episode()


if __name__ == "__main__":
    lerobot_builder(data_dir=os.path.join(os.path.dirname(__file__), '../data/modified_libero_rlds'), push_to_hub=False)